from .base import MetricBase
import torch
import numpy as np
import torch.nn.functional as F


class WMSE(MetricBase):
    def __init__(self, image_size = (64,64)):
        super().__init__()
        self.image_size = image_size
        
    def calculate_weights(self, y_true, eps = 1e-6):
        '''
        Calculate the weights of 0 pixels and 1 pixels,
        y_true here is 
        '''
        total_pixels = np.prod(self.image_size)
        freq_0 = (y_true == 0).sum(dim=(-3,-2,-1)) / total_pixels
        freq_1 = (y_true == 1).sum(dim=(-3,-2,-1)) / total_pixels
        freq_median = (freq_0 + freq_1) / 2
        w_0 = freq_median / (freq_0 + eps)
        w_1 = freq_median / (freq_1 + eps)
        return w_0, w_1
    
    def compute(self, y_true, y_pred):
        w_0, w_1 = self.calculate_weights(y_true)
        y_true = y_true.to(y_pred, non_blocking=True)
        mask = torch.where(y_true == 0, 
                           w_0.unsqueeze_(-1).unsqueeze_(-1).unsqueeze_(-1).expand(*y_true.size()), 
                           w_1.unsqueeze_(-1).unsqueeze_(-1).unsqueeze_(-1).expand(*y_true.size()))
        mse = F.mse_loss(y_pred, y_true, reduction='none')
        wmse = torch.sum(mask * mse) / y_true.numel()
        return wmse
    
    def summary(self):
        wmse = self.value / self.total
        print(f'WMSE: {wmse}')
        return wmse.item()

if __name__ == "__main__":
    wmse = WMSE(image_size=(64,64))
    y_true = torch.stack([torch.bernoulli(torch.full((64,64), 0.3)).unsqueeze(0) for _ in range(16)]).float()
    y_pred = torch.randint(0,2, size=(16,1,64,64)).float()
    print(y_true.shape)
    wmse(y_true, y_pred)
    wmse.summary()
    